""" Module responsible for downloading generif data from www and parse data
into a usable storage and query format.
"""
import pickle
import os
import sqlite3

import conf


class Fetch:
    """Base class for fetching (wget) data from www.

    Using wget *and gzip*, so its pretty much useless for windows.
    
    
    Parameters:
    -----------
    datadir : {str} directory for storage output
    url : {str} base url
    fnbase : {str} base filename (without postfix description)
    """
    def __init__(self, datadir=None, url='', fnbase=''):
        self.datadir = datadir
        self.url = url
        self.fnbase = fnbase
        self._data = None

    def load(self, fnbase):
        pass

    def save(self, fnbase):
        pass

    def _wget(self, filename):
        """ Wgets file from www. filename is relative."""
        if not self.url.endswith('/'):
            self.url += '/'
        full_path = self.url + filename
        if self.datadir == None:
            self.datadir = os.path.abspath('.')
        wd = os.path.abspath('.')
        call_str = "wget %s" % full_path
        os.chdir(self.datadir)
        print "Downloading file: %s" %full_path
        os.system(call_str)
        if filename.endswith('.gz'):
            print "Unzipping: %s ..." % filename
            os.system('gunzip %s' % filename)
        os.chdir(wd)
        print "Done."


class GeneRIF(Fetch):
    """Class responsible for storing and retrieving generif data."""
    def __init__(self, datadir=None, url=None, fnbase=None):
        self.connection = None
        if datadir == None:
            if hasattr(conf, "datadir"):
                datadir = conf.datadir 
            else:
                raise ValueError("Datadir needs to be specified.")
        if url == None:
            if hasattr(conf, "url"):
                url = conf.url 
            else:
                raise ValueError("url needs to be specified.")
        if fnbase == None:
            if hasattr(conf, "fnbase"):
                fnbase = conf.fnbase 
            else:
                raise ValueError("fnbase (filename) needs to be specified.")

        Fetch.__init__(self, datadir, url, fnbase)

        if self._check_update():
            print "An updated version of generifs is available"
        self.load()
        
    def load(self, fnbase='generifs_basic.pickle'):
        """ Loads pickled organism specific generif data."""
        try:
            fnbase = os.path.join(self.datadir, fnbase)
            fh = open(fnbase, "r")
        except:
            print "Could not open file: %s " % fnbase
            return None
        self._data = pickle.load(fh)
        print "%s loaded \n" % fnbase

    def save(self, fnbase='generifs_basic.pickle'):
        """Saves organsim specific generif data."""
        if self._data == None:
            print "No data to save"
            return
        fnbase = os.path.join(self.datadir, fnbase)
        fh = open(fnbase, "w")
        pickle.dump(self._data, fh, pickle.HIGHEST_PROTOCOL)
        print "Data saved to: %s" % fnbase

    # Private interface
    def _check_update(self):
        # fixme: not currently used
        
        # check if update is needed
        if not self._data:
            return True
        # check datetime of file and last download
        raise NotImplementedError

    def _text_parse_generifs(self, fnbase, org='9606'):
        """Parser for generif textfile at ncbi."""
        filename = os.path.join(self.datadir, fnbase+'.gz')
        if not os.path.exists(filename):
            filename = os.path.join(self.datadir, fnbase)
            if not os.path.exist(filename):
                raise IOError
        raw_text = open(filename).read().splitlines()
        grifs = defaultdict(set)
        for line in raw_text:
            els = line.split("\t")
            tax_id, gene, pmids, rif = els[0], els[1], els[2], els[-1]
            if org == None or org == tax_id: 
                grifs[geneid].add(rif)
                pmids[geneid].update(pmids.split(','))
        self._data = grifs
        # print "Done parsing generifs"
    
    def _sqlite_parse_generifs(self, fnbase):
        """Parse generif textfile into an sqlite database."""
        filename = os.path.join(self.datadir, fnbase)
        raw_text = open(filename).read().splitlines()
        for line in raw_text:
            els = line.split("\t")
        self.dbfile_abspath = os.path.abspath(self.datadir + 'grif.db')
        connection = sqlite3.connect(':memory:')
        cursor = connection.cursor()
        cursor.execute("""
        CREATE TABLE generifs
        (tax_id, gene_id, pmids, date, rif)
        """)
        # read generifs_basic
        rows = []
        for line in raw_text:
            rows.append(line.split("\t"))
        cursor.executemany("""
        INSERT INTO generifs
        (tax_id, gene_id, pmids, date, rif)
        values (?, ?, ?, ?,?)
        """, rows)
        self.connection = connection

    def _create_data(self, tax_id='9606', prim_key='gene_id'):
        # fixme: not used and needs work
        cursor = self.connection.cursor()
        cursor.execute("""
        SELECT gene_id, pmids, rif
        FROM generifs
        WHERE tax_id=(?)""", tax_id
                       )
        self._grifs = defaultdict(set)
        self._pmids = defaultdict(set)
        for row in cursor:
            self._grifs[row[0]].add(row[-1])
            self._pmids[row[0]].update(row[1].split(','))
        self._data = self._grifs
        print "Done"
